import os
import pandas as pd
from typing import Union, List
from tqdm import tqdm
from setting import raw_industry_data, address, StaticValues


def is_valid_address(
    file_name,
    prov: Union[List[str], str] = None,
    city: Union[List[str], str] = None,
) -> bool:
    """
    根据地址信息，筛选企业
    """

    df = pd.read_csv(file_name, nrows=1)
    try:
        if prov:
            if isinstance(prov, str):
                prov = [prov]

            sign = any(p in df.iloc[0]["所属省份"] for p in prov)
            if not sign:
                return False

        if city:
            if isinstance(city, str):
                city = [city]

            sign = any(c in df.iloc[0]["所属城市"] for c in city)
            if not sign:
                return False
    except:
        print(file_name)
        print(df.iloc[0])
    return True


tqdm.pandas()


def filter_by_keyword(
    df_: pd.DataFrame, cols: Union[List[str], str], keyword: Union[List[str], str]
) -> pd.DataFrame:
    """
    根据关键字筛选企业
    """
    if isinstance(cols, str):
        cols = [cols]
    if isinstance(keyword, str):
        keyword = [keyword]

    # 使用正则表达式模式匹配所有关键字（可选，但仅当关键字数量很大时可能更有效）
    # pattern = '|'.join(r'\b{}\b'.format(re.escape(kw)) for kw in keyword)

    # 直接在选定的列上应用str.contains，并使用any沿行方向聚合结果
    mask = df_[cols].progress_apply(
        lambda row: row.astype(str).str.contains("|".join(keyword), na=False).any(),
        axis=1,
    )
    return df_[mask]


# 示例用法
# df = pd.DataFrame({...})  # 您的DataFrame
# result = filter_by_keyword(df, 'column_name', 'keyword')
# 或者使用多个列和多个关键字
# result = filter_by_keyword(df, ['col1', 'col2'], ['kw1', 'kw2'])ries)


def get_df_by_kw_filter(address: dict, industry_kw: list):
    # 遍历 data_home_folder 下的所有文件，找出csv文件
    series = []
    for file_name in os.listdir(raw_industry_data):
        if file_name.endswith(".csv"):
            name = os.path.join(raw_industry_data, file_name)
            if is_valid_address(name, **address):
                # series.append(pd.read_csv(name, nrows=1))
                series.append(
                    filter_by_keyword(
                        pd.read_csv(name),
                        ["企业名称", "经营范围"],
                        industry_kw,
                    )
                )
                print("文件读取结束：", file_name)

    df = pd.concat(series)
    return df


def kw_cls(name, overwrite=False):
    global address
    sv = StaticValues(name)
    logger = sv.logger
    if os.path.exists(sv.KW_CSV):
        if not overwrite:
            logger.info(f"{sv.KW_CSV} 文件已存在，不再使用关键词过滤")
            return

    df = get_df_by_kw_filter(
        address=address,
        industry_kw=sv.KEY_WORDS,
    )
    df = df[df["经营状态"].isin(["存续", "开业", "在业", "正常"])]
    logger.info(f"关键词过滤后，共有 {len(df)} 家{sv.chinese_name}企业")
    df.to_csv(sv.KW_CSV, index=False)


if __name__ == "__main__":
    kw_cls("hydrogen")
